import sys, os
import tempfile
import shutil

sys.path.insert(1, os.path.join("..","..",".."))
from h2o.automl import H2OAutoML
from h2o.automl.autoh2o import get_automl
from tests import pyunit_utils as pu

from _automl_utils import import_dataset


def automl_checkpoints():
    ds = import_dataset()
    checkpoints_dir = tempfile.mkdtemp()

    aml = H2OAutoML(project_name="py_aml0", stopping_rounds=3, stopping_tolerance=0.001, stopping_metric="AUC",
                    max_models=2, seed=1234, export_checkpoints_dir=checkpoints_dir)
    aml.train(y=ds.target, training_frame=ds.train)

    get_aml = get_automl(aml.project_name)
    num_files = len([f for f in os.listdir(checkpoints_dir) if "_cv" not in f])  # do not count CV models
    shutil.rmtree(checkpoints_dir)

    assert aml.project_name == get_aml["project_name"]
    assert num_files > 0, "No models generated by AutoML"
    assert get_aml["leaderboard"].nrows == num_files, "Not all generated autoML models were saved."


pu.run_tests([
    automl_checkpoints
])
